from typing import List


class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        nums_sort = [0] * 10000
        sum_all = 0
        max_num = 0
        for num in nums:
            nums_sort[num] += 1
            sum_all += num
            max_num = max(max_num, num)
        avg_num = sum_all // k
        # 先处理特殊情况
        if sum_all % k != 0 or max_num > avg_num:
            return False

        # nums.sort()
        # print(nums)
        # print("avg: ", avg_num)
        # print("max: ", max_num)

        # 把数据放回去
        def push(n):
            max_num2 = max_num
            if n > max_num2:
                max_num2 = n
            nums_sort[n] += 1
            # print(max_num)
            return max_num2

        # 取一个数
        def pop(now_max_num, rest_num):
            max_num2 = max_num
            n = min(now_max_num, rest_num)
            for i in range(n, -1, -1):
                if nums_sort[i] != 0:
                    nums_sort[i] -= 1
                    if i == max_num2 and nums_sort[i] == 0:
                        # 寻找下一个最大值
                        for j in range(i, -1, -1):
                            # print("j:", j)
                            if nums_sort[j] != 0:
                                max_num2 = j
                                break
                    return i, max_num2
            # print(max_num)
            return -1, max_num2

        allSet = []
        while k > 0:
            # print("index: ", k)
            now_nums = []
            rest_num = avg_num
            now_max_num = max_num
            while True:
                # print("now_max_num:", now_max_num)
                # print("rest_num:", rest_num)
                n, max_num = pop(now_max_num, rest_num)
                # print("pop:", n)
                # print("now_nums:", now_nums)
                # print("now_nums:", now_nums)
                # 取不到合理的数据了
                if n == -1:
                    # 数组数据只剩一个了，就说明没有与这个数据匹配的了
                    if len(now_nums) == 1 and len(allSet) == 0 or len(now_nums) == 0:
                        return False
                    elif len(now_nums) == 1 and len(allSet) != 0:
                        # 回退上一个集合
                        last_num = now_nums.pop()
                        max_num = push(last_num)
                        now_nums = allSet.pop()
                        # 回退上一个的最后一个元素
                        last_num = now_nums.pop()
                        max_num = push(last_num)
                        now_max_num = last_num - 1
                        rest_num = avg_num - sum(now_nums)
                        k += 1
                        # print("回退一轮")
                        continue
                    else:
                        # 把数据放回去，取下一个
                        last_num = now_nums.pop()
                        max_num = push(last_num)
                        now_max_num = last_num - 1
                        rest_num += last_num
                        continue
                now_max_num = n
                rest_num = rest_num - n
                now_nums.append(n)
                if rest_num == 0:
                    # print(now_nums)
                    allSet.append(now_nums)
                    k -= 1
                    # print("\n\n")
                    break
        # max_num = push(i * 10, max_num)
        # print(max_num)

        return True


if __name__ == '__main__':
    # nums = [4, 5, 9, 3, 10, 2, 10, 7, 10, 8, 5, 9, 4, 6, 4, 9]
    # k = 5
    nums = [10, 1, 10, 9, 6, 1, 9, 5, 9, 10, 7, 8, 5, 2, 10, 8]
    k = 11
    s = Solution()

    print(s.canPartitionKSubsets(nums, k))

    # [2, 3, 4, 4, 4, 5, 5, 6, 7, 8, 9, 9, 9, 10, 10, 10]
    #
    # 10, 9, 2
    # 10, 8, 3
    # 10, 6, 5
    # [ 4, 4, 4, 5, 7, 9, 9]
    # 9, 7,5
    # 9 4 4 4
